import numpy as np


# def binarize_ret(df):
#     # 修正为 0-3 编码
#     conditions = [
#        (df['ret3'] > 0.15) ,   # →3
#         (df['ret3'] > 0.05) & (df['ret3'] <= 0.15),  # →2
#         (df['ret3'] > 0) & (df['ret3'] <= 0.05),     # →1
#         df['ret3'] <= 0             # →0
#     ]
#     choices = [3, 2, 1, 0]  # 标签范围调整为 0-3
#     df['label'] = np.select(conditions, choices, default=np.nan)
#     return df

def binarize_ret_10cm_gt4(df):
    # 初始化所有标签为0（默认值）
    df['label'] = 0
    
    # 优先级1: 满足任意一个高增长条件 → label=3
    cond3 = ((df['ret1'] > 0.04) | (df['ret2'] > 0.09) | (df['ret3'] > 0.21) )&((df['ret3'] > 0.05))
    df.loc[cond3, 'label'] = 3
    
    # 优先级2: 满足任意一个中等增长条件 → label=2 (且未标记为3)
    cond2 = (df['ret1'] > 0.01) | (df['ret2'] > 0.05) |  (df['ret3'] > 0.10)
    df.loc[cond2 & (df['label'] == 0), 'label'] = 2  # 仅更新未标记的样本
    
    # 优先级3: 满足任意一个低增长条件 → label=1 (且未标记为2或3)
    cond1 = (df['ret1'] > 0) | (df['ret2'] > 0) |  (df['ret3'] > 0)
    df.loc[cond1 & (df['label'] == 0), 'label'] = 1  # 仅更新未标记的样本
    
    # 优先级4: 所有收益率都小于0 → 保持label=0
    # 不需要额外操作，因为默认已经是0
    return df
    

def binarize_ret_10cm_zt(df):
    # 初始化所有标签为0（默认值）
    df['label'] = 0
    
    # 优先级1: 满足任意一个高增长条件 → label=3
    cond3 = (df['ret2'] > 0.12) | (df['ret3'] > 0.21)  
    df.loc[cond3, 'label'] = 3
    
    # 优先级2: 满足任意一个中等增长条件 → label=2 (且未标记为3)
    cond2 = (df['ret1'] > 0.04) | (df['ret2'] > 0.05) |  (df['ret3'] > 0.10)
    df.loc[cond2 & (df['label'] == 0), 'label'] = 2  # 仅更新未标记的样本
    
    # 优先级3: 满足任意一个低增长条件 → label=1 (且未标记为2或3)
    cond1 = (df['ret1'] > 0) | (df['ret2'] > 0) |  (df['ret3'] > 0)
    df.loc[cond1 & (df['label'] == 0), 'label'] = 1  # 仅更新未标记的样本
    
    # 优先级4: 所有收益率都小于0 → 保持label=0
    # 不需要额外操作，因为默认已经是0
    
    return df

    
    
def binarize_ret_w_ma(df):
    # 修正为 0-3 编码
    conditions = [
        df['ret3'] > 0.18,          # →3
        (df['ret3'] > 0.10) & (df['ret3'] <= 0.18),  # →2
        (df['ret3'] > 0.05) & (df['ret3'] <= 0.10),     # →1
        df['ret3'] <= 0.05             # →0
    ]
    choices = [3, 2, 1, 0]  # 标签范围调整为 0-3
    df['label'] = np.select(conditions, choices, default=np.nan)
    return df

def binarize_ret_20cm(df):
    # 修正为 0-3 编码
    conditions = [
        df['ret3'] > 0.15,          # →3
        (df['ret3'] > 0.05) & (df['ret3'] <= 0.15),  # →2
        (df['ret3'] > 0) & (df['ret3'] <= 0.05),     # →1
        df['ret3'] <= 0             # →0
    ]
    choices = [3, 2, 1, 0]  # 标签范围调整为 0-3
    df['label'] = np.select(conditions, choices, default=np.nan)
    return df


# def binarize_ret_20cm(df):
#     # 修正为 0-3 编码
#     conditions = [
#         (df['ret3'] > 0.15) ,  # →2
#         (df['ret3'] > 0.05) & (df['ret3'] <= 0.15),     # →1
#         df['ret3'] <= 0.05             # →0
#     ]
#     choices = [2, 1, 0]  # 标签范围调整为 0-3
#     df['label'] = np.select(conditions, choices, default=np.nan)
#     return df